from torch import nn
import torch
import torch.nn.functional as F


class VGG16(nn.Module):
    def __init__(self, num_classes=10):  # 设置类别为10。CIFAR10数据集共有10个类别
        super(VGG16, self).__init__()
        self.block_1 = nn.Sequential(  # 第一个卷积块，两个卷积层，一个最大池化层
            nn.Conv2d(3, 64, 3, 1, 1),  # 进行卷积运算。输入通道3，输出通道64，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活，inplace=True表示覆盖原来的变量名称，这样可以节省内存
            nn.Conv2d(64, 64, 3, 1, 1),  # 进行卷积运算。输入通道64，输出通道64，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),  # 最大池化，核大小为2*2，步长为2，不填充
        )
        self.block_2 = nn.Sequential(  # 第二个卷积块，两个卷积层，一个最大池化层
            nn.Conv2d(64, 128, 3, 1, 1),  # 进行卷积运算。输入通道64，输出通道128，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.Conv2d(128, 128, 3, 1, 1),  # 进行卷积运算。输入通道128，输出通道128，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.MaxPool2d(2, 2, 0),  # 最大池化，核大小为2*2，步长为2，不填充
        )
        self.block_3 = nn.Sequential(  # 第三个卷积块，三个卷积层，一个最大池化层
            nn.Conv2d(128, 256, 3, 1, 1),  # 进行卷积运算。输入通道128，输出通道256，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.Conv2d(256, 256, 3, 1, 1),  # 进行卷积运算。输入通道256，输出通道256，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.Conv2d(256, 256, 3, 1, 1),  # 进行卷积运算。输入通道256，输出通道256，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.MaxPool2d(2, 2, 0),  # 最大池化，核大小为2*2，步长为2，不填充
        )
        self.block_4 = nn.Sequential(  # 第四个卷积块，三个卷积层，一个最大池化层
            nn.Conv2d(256, 512, 3, 1, 1),  # 进行卷积运算。输入通道256，输出通道512，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.Conv2d(512, 512, 3, 1, 1),  # 进行卷积运算。输入通道512，输出通道512，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.Conv2d(512, 512, 3, 1, 1),  # 进行卷积运算。输入通道512，输出通道512，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.MaxPool2d(2, 2, 0),  # 最大池化，核大小为2*2，步长为2，不填充
        )
        self.block_5 = nn.Sequential(  # 第五个卷积块，三个卷积层，一个最大池化层
            nn.Conv2d(512, 512, 3, 1, 1),  # 进行卷积运算。输入通道512，输出通道512，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.Conv2d(512, 512, 3, 1, 1),  # 进行卷积运算。输入通道512，输出通道512，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.Conv2d(512, 512, 3, 1, 1),  # 进行卷积运算。输入通道512，输出通道512，卷积核3*3，步长1，填充1
            nn.ReLU(inplace=True),  # relu激活
            nn.MaxPool2d(2, 2, 0),  # 最大池化，核大小为2*2，步长为2，不填充
        )
        self.classifier = nn.Sequential(  # 全连接层
            nn.Flatten(),  # 将上一步的结果进行压平
            nn.Linear(512, 256),  # 线性运算，输入512个元素，输出256个
            nn.ReLU(inplace=True),  # relu激活
            nn.Linear(256, 256),  # 线性运算，输入256个元素，输出256
            nn.ReLU(inplace=True),  # relu激活
            nn.Linear(256, num_classes),  # 线性运算，输入256，输出10个类别
        )

    def forward(self, x):  # 前向传播
        x = self.block_1(x)  # 不断地调用上述的网络进行训练
        x = self.block_2(x)
        x = self.block_3(x)
        x = self.block_4(x)
        x = self.block_5(x)
        # x = self.classifier(x.view(-1, 512 * 1 * 1))  # x.view(-1,512*1*1)也可以进行压平，-1表示让计算机自己计算batch_size的大小
        x = self.classifier(x)
        x = F.softmax(x, dim=1)   # 最后使用softmax函数进行激活
        return x


if __name__ == '__main__':  # 测试模型是否正确
    model = VGG16()
    input = torch.randn((1, 3, 32, 32))
    output = model(input)
    print(output.shape)
